[BUG] Use models literal StructuredDataset to enable sd bypass task #2954

JiangJiaWei1103 · 2024-11-25T15:31:28Z

Tracking issue

Why are the changes needed?

For a flyte task with python StructuredDataset as an input from dataclass attribute access (i.e., dc.sd), it fails to convert the task output LiteralMap back to flyte idl. On the highest level, the process of calling _dispatch_execute (here) is illustrated as follows:

As can be observed, the problem (marked as orange) occurs when the input LiteralMap is converted to python natives, during which a Literal is built with python native StructuredDataset, instead of using the correct literals.StructuredDataset (please refer to this code snippet).

This makes conversion from the task output LiteralMap back to flyte idl fail because python native StructuredDataset has no to_flyte_idl method.

What changes were proposed in this pull request?

Directly construct a literals.StructuredDataset within dict_to_structured_dataset method of StructuredDatasetTransformerEngine.

How was this patch tested?

Execute the command run in the K8s pod, as suggested by @pingsutw:
- Describe the pod executing the remote task in the reported issue.
- Copy the command in the pod and reproduce locally.

After fixing, outputs.pb now can be successfully written to s3 storage.

Do remote sandbox test with flytekit installed by commit sha in this PR.

from dataclasses import dataclass, field

from flytekit import task, workflow, ImageSpec
from flytekit.types.structured import StructuredDataset


# Build docker image and push to the local registry
flytekit_hash = "d5d8620cc0494a148adb6bb4a13978d810d9673a"  # Refer to this PR
flytekit = f"git+https://github.com/flyteorg/flytekit.git@{flytekit_hash}"
image = ImageSpec(
    packages=[
        flytekit,
        "pandas",
        "pyarrow"
    ],
    apt_packages=["git"],
    registry="localhost:30000",
)


@dataclass
class DC:
    sd: StructuredDataset = field(
        default_factory=lambda: StructuredDataset(
            uri="s3://my-s3-bucket/s3_flyte_dir/df.parquet", 
            file_format="parquet"
        )
    )


@task(container_image=image)
def t_sd_attr(sd: StructuredDataset) -> StructuredDataset:
    return sd


@workflow
def wf(dc: DC):
    t_sd_attr(sd=dc.sd)

Run the following command:

pyflyte run --remote reprod.py wf --dc '{"dc": {"sd": {"uri": "s3://my-s3-bucket/s3_flyte_dir/df.parquet", "file_format": "parquet"}}}'

The remote sandbox test result is shown as follows:

Setup process

git clone https://github.com/flyteorg/flytekit.git
gh pr checkout 2954
make setup && pip install -e .

Check all the applicable boxes

I updated the documentation accordingly.
All new and existing tests passed.
All commits are signed-off.

Related PRs

#2914

Docs link

Signed-off-by: JiaWei Jiang <[email protected]>

JiangJiaWei1103 · 2024-11-26T14:29:29Z

Follow-up

Because the reported bug is raised in the remote run (which involves communication with s3), I think it's hard to test in both unit test and integration test, as we've discussed before here. If there's any possibility to better test this fix, I'll handle it!

Future-Outlier · 2024-11-26T15:06:26Z

flytekit/types/structured/structured_dataset.py

@@ -197,12 +197,25 @@ def return_sd() -> StructuredDataset:
            return df

        For details, please refer to this issue: https://github.com/flyteorg/flyte/issues/5954.
+
+        2. Need access to self._literal_sd when converting task output back to flyteidl, please see:
+        https://github.com/flyteorg/flytekit/blob/master/flytekit/bin/entrypoint.py#L326


Can we use github permanent link?

Sure, let me fix it. Thanks!

Signed-off-by: JiaWei Jiang <[email protected]>

JiangJiaWei1103 · 2024-11-28T16:23:53Z

Follow-up

The last commit aims at simplifying the building logic of sd_literal. Instead of constructing a dummy python native StructuredDataset then calling set_literal, we directly build sd_literal here. Thanks for @Future-Outlier 's suggestion.

The next steps are:

Make sure this patch passes the sandbox test using the image built with the last commit sha.
Handle the empty-string file_format issue.

Future-Outlier · 2024-11-28T16:25:17Z

Follow-up

The last commit aims at simplifying the building logic of sd_literal. Instead of constructing a dummy python native StructuredDataset then calling set_literal, we directly build sd_literal here. Thanks for @Future-Outlier 's suggestion.

The next steps are:

Make sure this patch passes the sandbox test using the image built with the last commit sha.

Handle the empty-string file_format issue.

you can use single binary to test it, this will be quicker.
this will help a lot, thank you man

codecov · 2024-11-28T16:36:03Z

Codecov Report

Attention: Patch coverage is 11.11111% with 8 lines in your changes missing coverage. Please review.

Project coverage is 78.28%. Comparing base (f938661) to head (7ce7ea8).
Report is 1 commits behind head on master.

Files with missing lines	Patch %	Lines
flytekit/types/structured/structured_dataset.py	11.11%	8 Missing ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##           master    #2954      +/-   ##
==========================================
+ Coverage   75.14%   78.28%   +3.13%     
==========================================
  Files         200      200              
  Lines       20919    20928       +9     
  Branches     2692     2693       +1     
==========================================
+ Hits        15720    16383     +663     
+ Misses       4445     3745     -700     
- Partials      754      800      +46

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

JiangJiaWei1103 · 2024-11-29T15:29:49Z

Hi @Future-Outlier,

The last commit passes the test example you provided me yesterday, as shown below:

For the first question, I do use a single binary to run remote tests as suggested in the official setup. In step 2, flytectl demo start --dev is used to start the k3s cluster and this environment is referred to as Flyte sandbox, so I just call it a sandbox test. Sorry for the dumb question here, but I would like to further clarify the difference and refine my understanding. Thanks very much!

For the empty-string file_format, I suspect this information is erased before inputs.pb is loaded from s3. Following shows the loaded input proto:

As can be observed, the information of file_format is absent. Hence, I think I'll continue tracing the serialization and registration process to find out what's going on.

Signed-off-by: JiaWei Jiang <[email protected]>

JiangJiaWei1103 · 2024-11-30T08:54:48Z

Further Discussion - Logic behind `extract_cols_and_format`

For the last commit, we drop the redundant handling of column information of StructuredDataset in dict_to_structured_dataset. First of all, extract_cols_and_format isn't exclusive to deprecated FlyteSchema. It's used to iterate through Annotated and extract the information like base type, column information, etc. Following demonstrates a simple example about extracting one column out of a pd.DataFrame and returning one-column df:

all_cols = kwtypes(name=str, height=int)
col = kwtypes(height=int)

@task(container_image=image)
def t1(sd: Annotated[StructuredDataset, all_cols]) -> Annotated[StructuredDataset, col]:
    df = sd.open(pd.DataFrame).all()

    return StructuredDataset(dataframe=df)

For these cases, async_to_python_value can already handle column information, as can be seen here and here. Hence, we just drop the redundant handling.

Test Result

Due to the limitation on unit tests and integration tests now, we only test this patch by single binary.

Empty-String `file_format`

I'll open another issue containing complete experiments and observations this week. And, I'll continue solving that problem.

Future-Outlier

Really Nice PR.
I can tell you put lots of efforts.
Let's make Flyte GREAT!

JiangJiaWei1103 · 2024-12-02T11:32:26Z

Thanks for your support and patience, bro!

Let's move on and make Flyte better!!

…lyteorg#2954) * fix: Use models literal sd to enable sd bypass task Signed-off-by: JiaWei Jiang <[email protected]> * Reuse sd _set_literal with a public wrapper method Signed-off-by: JiaWei Jiang <[email protected]> * Use github permanent link Signed-off-by: JiaWei Jiang <[email protected]> * Simplify sd_literal building logic Signed-off-by: JiaWei Jiang <[email protected]> * Drop unnecessary column handling Signed-off-by: JiaWei Jiang <[email protected]> --------- Signed-off-by: JiaWei Jiang <[email protected]> Co-authored-by: Future-Outlier <[email protected]>

JiangJiaWei1103 added 3 commits November 25, 2024 23:20

fix: Use models literal sd to enable sd bypass task

4595d46

Signed-off-by: JiaWei Jiang <[email protected]>

Merge branch 'master' into fix-sd-bypass-task

6bf3dd0

Reuse sd _set_literal with a public wrapper method

d5d8620

Signed-off-by: JiaWei Jiang <[email protected]>

JiangJiaWei1103 marked this pull request as ready for review November 26, 2024 14:15

JiangJiaWei1103 requested review from wild-endeavor, kumare3, eapolinario, pingsutw, cosmicBboy, samhita-alla, thomasjpfan and Future-Outlier as code owners November 26, 2024 14:15

Future-Outlier reviewed Nov 26, 2024

View reviewed changes

JiangJiaWei1103 added 2 commits November 26, 2024 23:22

Use github permanent link

9182b96

Signed-off-by: JiaWei Jiang <[email protected]>

Simplify sd_literal building logic

7ce7ea8

Signed-off-by: JiaWei Jiang <[email protected]>

Drop unnecessary column handling

adc1061

Signed-off-by: JiaWei Jiang <[email protected]>

JiangJiaWei1103 mentioned this pull request Nov 30, 2024

[test] Add integration test for accessing sd sttr in dc #2969

Merged

3 tasks

Merge branch 'master' into fix-sd-bypass-task

4e467b9

Future-Outlier approved these changes Dec 2, 2024

View reviewed changes

Future-Outlier enabled auto-merge (squash) December 2, 2024 03:21

Future-Outlier merged commit 0b0e338 into flyteorg:master Dec 2, 2024
29 checks passed

Future-Outlier mentioned this pull request Feb 22, 2025

Nominate @jiangjiawei1103 as Committer flyteorg/community#37

Closed

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

[BUG] Use models literal StructuredDataset to enable sd bypass task #2954

[BUG] Use models literal StructuredDataset to enable sd bypass task #2954

JiangJiaWei1103 commented Nov 25, 2024 •

edited

Loading

JiangJiaWei1103 commented Nov 26, 2024 •

edited

Loading

Future-Outlier Nov 26, 2024

JiangJiaWei1103 Nov 26, 2024

JiangJiaWei1103 commented Nov 28, 2024 •

edited

Loading

Future-Outlier commented Nov 28, 2024

Follow-up

codecov bot commented Nov 28, 2024 •

edited

Loading

JiangJiaWei1103 commented Nov 29, 2024

JiangJiaWei1103 commented Nov 30, 2024 •

edited

Loading

Future-Outlier left a comment

JiangJiaWei1103 commented Dec 2, 2024

[BUG] Use models literal StructuredDataset to enable sd bypass task #2954

[BUG] Use models literal StructuredDataset to enable sd bypass task #2954

Conversation

JiangJiaWei1103 commented Nov 25, 2024 • edited Loading

Tracking issue

Why are the changes needed?

What changes were proposed in this pull request?

How was this patch tested?

Setup process

Check all the applicable boxes

Related PRs

Docs link

JiangJiaWei1103 commented Nov 26, 2024 • edited Loading

Follow-up

Future-Outlier Nov 26, 2024

Choose a reason for hiding this comment

JiangJiaWei1103 Nov 26, 2024

Choose a reason for hiding this comment

JiangJiaWei1103 commented Nov 28, 2024 • edited Loading

Follow-up

Future-Outlier commented Nov 28, 2024

Follow-up

codecov bot commented Nov 28, 2024 • edited Loading

Codecov Report

JiangJiaWei1103 commented Nov 29, 2024

JiangJiaWei1103 commented Nov 30, 2024 • edited Loading

Further Discussion - Logic behind extract_cols_and_format

Test Result

Empty-String file_format

Future-Outlier left a comment

Choose a reason for hiding this comment

JiangJiaWei1103 commented Dec 2, 2024

JiangJiaWei1103 commented Nov 25, 2024 •

edited

Loading

JiangJiaWei1103 commented Nov 26, 2024 •

edited

Loading

JiangJiaWei1103 commented Nov 28, 2024 •

edited

Loading

codecov bot commented Nov 28, 2024 •

edited

Loading

JiangJiaWei1103 commented Nov 30, 2024 •

edited

Loading

Further Discussion - Logic behind `extract_cols_and_format`

Empty-String `file_format`